
import random
import networkx as nx
from cluster_qkd_networkx_cluster import absolute_path

def getGraphWithTemplateAndCapacity(sinle_cluster_size: int, cluster_number: int, template_file_names: list, capacity: int) -> nx.Graph:
    # template_file_names指定了读入簇的文件名，单簇的节点数量应与sinle_cluster_size匹配，簇的个数应与cluster_number匹配
    # 返回一个网络
    graph = nx.Graph()
    for i in range(cluster_number):
        template_file_name = template_file_names[i]
        edges = extract_edges_from_file(template_file_name)
        for edge in edges:
            graph.add_edge(edge[0]+i*sinle_cluster_size, edge[1]+i*sinle_cluster_size, weight=edge[2])
    
    # 随机连簇内
    for j in range(cluster_number-1):
        for k in (range(j+1, cluster_number)):
            for _ in range(5):
                across_cluster_start_node = random.randint(1+j*sinle_cluster_size, sinle_cluster_size*(1+j))
                across_cluster_end_node = random.randint(1+k*sinle_cluster_size, sinle_cluster_size*(1+k))
                graph.add_edge(across_cluster_start_node, across_cluster_end_node, weight=random.randint(5, capacity*2))

    return graph

def getGraphWithTemplate(sinle_cluster_size: int, cluster_number: int, template_file_names: list) -> nx.Graph:
    # template_file_names指定了读入簇的文件名，单簇的节点数量应与sinle_cluster_size匹配，簇的个数应与cluster_number匹配
    # 返回一个网络
    graph = nx.Graph()
    for i in range(cluster_number):
        template_file_name = template_file_names[i]
        edges = extract_edges_from_file(template_file_name)
        for edge in edges:
            graph.add_edge(edge[0]+i*sinle_cluster_size, edge[1]+i*sinle_cluster_size, weight=edge[2])
    
    # 随机连簇内
    for j in range(cluster_number-1):
        for k in (range(j+1, cluster_number)):
            for _ in range(5):
                across_cluster_start_node = random.randint(1+j*sinle_cluster_size, sinle_cluster_size*(1+j))
                across_cluster_end_node = random.randint(1+k*sinle_cluster_size, sinle_cluster_size*(1+k))
                graph.add_edge(across_cluster_start_node, across_cluster_end_node, weight=random.randint(5, 100))

    return graph


def extract_edges_from_file(filename):
    with open(f"{absolute_path}template/{filename}", 'r') as file:
        lines = file.readlines()
    num_nodes = int(lines[0].strip())
    network_params = lines[1].strip().split()
    num_edges = int(lines[2].strip())
    edges = []
    for line in lines[3:3+num_edges]:
        edge_info = list(map(int, line.strip().split()))
        edges.append(edge_info)
    return edges

